import akshare as ak
import pandas as pd
import datetime
from tabulate import tabulate
import talib
import numpy as np
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

def test_data_availability(stock_code="600519"):
    """测试各类数据的可获取性"""
    logger.info(f"开始测试股票 {stock_code} 的数据可获取性")
    results = []
    end_date = datetime.datetime.now().strftime('%Y%m%d')
    start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y%m%d')

    # 基础价格数据测试
    try:
        daily_data = ak.stock_zh_a_hist(symbol=stock_code, start_date=start_date, end_date=end_date)
        logger.info(f"获取到 {len(daily_data)} 条K线数据")
        results.append(["基础价格数据", "K线数据", "✓", "可获取日K线数据"])
    except Exception as e:
        logger.error(f"获取K线数据失败: {str(e)}")
        results.append(["基础价格数据", "K线数据", "✗", "获取失败"])
    
    try:
        min_data = ak.stock_zh_a_hist_min_em(symbol=stock_code)
        logger.info(f"获取到 {len(min_data)} 条分时数据")
        results.append(["基础价格数据", "分时数据", "✓", "可获取分时数据"])
    except Exception as e:
        logger.error(f"获取分时数据失败: {str(e)}")
        results.append(["基础价格数据", "分时数据", "✗", "获取失败"])

    # 技术指标数据测试
    try:
        macd_data = ak.stock_zh_a_hist(symbol=stock_code, start_date=start_date, end_date=end_date)
        logger.info("成功获取MACD计算所需数据")
        results.append(["技术指标数据", "MACD", "✓", "可通过K线数据计算"])
    except Exception as e:
        logger.error(f"获取MACD数据失败: {str(e)}")
        results.append(["技术指标数据", "MACD", "✗", "获取失败"])

    # 资金流向数据测试
    try:
        # 根据股票代码判断市场类型
        market = "sh" if stock_code.startswith('6') else "sz" if stock_code.startswith(('0', '3')) else "bj"
        flow_data = ak.stock_individual_fund_flow(stock=stock_code, market=market)
        if flow_data is None or flow_data.empty:
            logger.warning(f"获取股票 {stock_code} 的资金流向数据失败")
            results.append(["资金流向数据", "主力资金流向", "✗", "数据为空"])
        else:
            logger.info(f"获取到 {len(flow_data)} 条资金流向数据")
            results.append(["资金流向数据", "主力资金流向", "✓", "可获取资金流向数据"])
    except Exception as e:
        logger.error(f"获取资金流向数据失败: {str(e)}")
        results.append(["资金流向数据", "主力资金流向", "✗", "获取失败"])
    
    try:
        north_data = ak.stock_hsgt_north_net_flow_in_em()
        logger.info(f"获取到 {len(north_data)} 条北向资金数据")
        results.append(["资金流向数据", "北向资金", "✓", "可获取北向资金数据"])
    except Exception as e:
        logger.error(f"获取北向资金数据失败: {str(e)}")
        results.append(["资金流向数据", "北向资金", "✗", "获取失败"])

    # 市场情绪数据测试
    try:
        sentiment_data = ak.stock_market_activity_legu()
        logger.info(f"获取到 {len(sentiment_data)} 条市场活跃度数据")
        results.append(["市场情绪数据", "市场活跃度", "✓", "可获取市场活跃度数据"])
    except Exception as e:
        logger.error(f"获取市场活跃度数据失败: {str(e)}")
        results.append(["市场情绪数据", "市场活跃度", "✗", "获取失败"])

    # 基本面数据测试
    try:
        financial_data = ak.stock_financial_report_sina(stock=stock_code)
        logger.info(f"获取到 {len(financial_data)} 条财务数据")
        results.append(["基本面数据", "财务数据", "✓", "可获取财务报告数据"])
    except Exception as e:
        logger.error(f"获取财务数据失败: {str(e)}")
        results.append(["基本面数据", "财务数据", "✗", "获取失败"])
    
    try:
        valuation_data = ak.stock_a_lg_indicator(symbol=stock_code)
        logger.info(f"获取到 {len(valuation_data)} 条估值数据")
        results.append(["基本面数据", "估值数据", "✓", "可获取估值指标数据"])
    except Exception as e:
        logger.error(f"获取估值数据失败: {str(e)}")
        results.append(["基本面数据", "估值数据", "✗", "获取失败"])

    # 消息面数据测试
    try:
        news_data = ak.stock_news_em(symbol=stock_code)
        logger.info(f"获取到 {len(news_data)} 条新闻数据")
        results.append(["消息面数据", "公司新闻", "✓", "可获取公司新闻数据"])
    except Exception as e:
        logger.error(f"获取新闻数据失败: {str(e)}")
        results.append(["消息面数据", "公司新闻", "✗", "获取失败"])

    # 大盘/行业对比数据测试
    try:
        index_data = ak.stock_zh_index_daily(symbol="sh000001")
        logger.info(f"获取到 {len(index_data)} 条指数数据")
        results.append(["大盘/行业对比数据", "指数数据", "✓", "可获取指数数据"])
    except Exception as e:
        logger.error(f"获取指数数据失败: {str(e)}")
        results.append(["大盘/行业对比数据", "指数数据", "✗", "获取失败"])

    # 回测特有数据测试
    try:
        dividend_data = ak.stock_history_dividend()
        logger.info(f"获取到 {len(dividend_data)} 条分红数据")
        results.append(["回测特有数据", "历史分红数据", "✓", "可获取分红数据"])
    except Exception as e:
        logger.error(f"获取分红数据失败: {str(e)}")
        results.append(["回测特有数据", "历史分红数据", "✗", "获取失败"])

    # 打印结果表格
    logger.info("\n数据可获取性测试结果：")
    print(tabulate(results, headers=["数据类别", "指标名称", "是否可获取", "备注"], tablefmt="grid"))

def get_specific_date_data(stock_code="600519", date="20240228"):
    """获取指定日期的K线和分时数据"""
    print(f"\n获取{stock_code}在{date}的数据：")
    
    # 获取K线数据
    try:
        daily_data = ak.stock_zh_a_hist(symbol=stock_code, start_date=date, end_date=date)
        if not daily_data.empty:
            print("\nK线数据：")
            print(tabulate(daily_data, headers="keys", tablefmt="grid", showindex=False))
        else:
            print("\nK线数据：无数据")
    except Exception as e:
        print(f"\nK线数据获取失败：{str(e)}")
    
    # 获取分时数据
    try:
        min_data = ak.stock_zh_a_hist_min_em(symbol=stock_code)
        if not min_data.empty:
            # 过滤出指定日期的分时数据
            min_data['日期'] = pd.to_datetime(min_data['时间']).dt.strftime('%Y%m%d')
            min_data = min_data[min_data['日期'] == date]
            if not min_data.empty:
                print("\n分时数据：")
                print(tabulate(min_data, headers="keys", tablefmt="grid", showindex=False))
            else:
                print("\n分时数据：该日期无数据")
        else:
            print("\n分时数据：无数据")
    except Exception as e:
        print(f"\n分时数据获取失败：{str(e)}")

def get_technical_indicators(stock_code="600519", date="20240228"):
    """获取指定日期的技术指标数据"""
    print(f"\n获取{stock_code}在{date}的技术指标数据：")
    
    # 获取前60天的数据用于计算指标
    start_date = (datetime.datetime.strptime(date, '%Y%m%d') - datetime.timedelta(days=60)).strftime('%Y%m%d')
    end_date = date
    
    try:
        # 获取日K线数据
        df = ak.stock_zh_a_hist(symbol=stock_code, start_date=start_date, end_date=end_date, adjust="qfq")
        if df.empty:
            print("\n无法获取K线数据，可能是非交易日或数据不可用")
            return
        
        # 转换数据类型
        close = df['收盘'].astype(float).values
        high = df['最高'].astype(float).values
        low = df['最低'].astype(float).values
        volume = df['成交量'].astype(float).values
        dates = df['日期'].values
        
        # 计算并打印每个交易日的技术指标
        print("\n每个交易日的技术指标数据：")
        for i in range(len(close)):
            print(f"\n交易日期：{dates[i]}")
            indicators = []
            
            # MACD
            macd, macdsignal, macdhist = talib.MACD(close[:i+1], fastperiod=12, slowperiod=26, signalperiod=9)
            indicators.append(["MACD", 
                             f"DIFF: {round(macd[-1], 2)}", 
                             f"DEA: {round(macdsignal[-1], 2)}", 
                             f"柱: {round(macdhist[-1], 2)}"])
            
            # RSI
            rsi6 = talib.RSI(close[:i+1], timeperiod=6)
            rsi12 = talib.RSI(close[:i+1], timeperiod=12)
            rsi24 = talib.RSI(close[:i+1], timeperiod=24)
            indicators.append(["RSI", 
                             f"RSI6: {round(rsi6[-1], 2)}", 
                             f"RSI12: {round(rsi12[-1], 2)}", 
                             f"RSI24: {round(rsi24[-1], 2)}"])
            
            # KDJ
            k, d = talib.STOCH(high[:i+1], low[:i+1], close[:i+1])
            j = 3 * k - 2 * d
            indicators.append(["KDJ", 
                             f"K: {round(k[-1], 2)}", 
                             f"D: {round(d[-1], 2)}", 
                             f"J: {round(j[-1], 2)}"])
            
            # 布林带
            upper, middle, lower = talib.BBANDS(close[:i+1], timeperiod=20, nbdevup=2.0, nbdevdn=2.0, matype=0)
            indicators.append(["BOLL", 
                             f"上轨: {round(upper[-1], 2)}", 
                             f"中轨: {round(middle[-1], 2)}", 
                             f"下轨: {round(lower[-1], 2)}"])
            
            # 成交量指标
            if i >= 4:  # 确保有足够数据计算MA5
                volume_ma5 = talib.MA(volume[:i+1], timeperiod=5)
                volume_ma10 = talib.MA(volume[:i+1], timeperiod=10) if i >= 9 else [0]
                indicators.append(["Volume MA", 
                                 f"MA5: {round(volume_ma5[-1], 2)}", 
                                 f"MA10: {round(volume_ma10[-1], 2)}", 
                                 f"成交量: {round(volume[i], 2)}"])
            
            print(tabulate(indicators, headers=["指标", "参数1", "参数2", "参数3"], tablefmt="grid"))
            print("-" * 80)  # 分隔线
        
        # 获取资金流向数据
        try:
            # 根据股票代码判断市场类型
            market = "sh" if stock_code.startswith('6') else "sz" if stock_code.startswith(('0', '3')) else "bj"
            flow_data = ak.stock_individual_fund_flow(stock=stock_code, market=market)
            if not flow_data.empty:
                flow_data['日期'] = pd.to_datetime(flow_data['日期']).dt.strftime('%Y%m%d')
                flow_data = flow_data[flow_data['日期'] == date]
                if not flow_data.empty:
                    print("\n资金流向数据：")
                    print(tabulate(flow_data, headers="keys", tablefmt="grid", showindex=False))
                else:
                    print("\n资金流向数据：该日期无数据")
        except Exception as e:
            print(f"\n资金流向数据获取失败：{str(e)}")
    
    except Exception as e:
        print(f"\n技术指标数据获取失败：{str(e)}")

if __name__ == "__main__":
    # 获取命令行参数
    import sys
    if len(sys.argv) > 1:
        stock_code = sys.argv[1]
        # 如果提供了日期参数，则运行技术指标测试
        if len(sys.argv) > 2:
            date = sys.argv[2]
            get_technical_indicators(stock_code, date)
        else:
            # 否则运行完整的数据可用性测试
            test_data_availability(stock_code)
    else:
        # 默认使用比亚迪作为示例，运行完整的数据可用性测试
        test_data_availability("002594")